package com.shujia.spark.sql

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.{DataFrame, SparkSession}

object Demo5Window {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("window")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val studentDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("id STRING,name STRING,age INT,sex STRING,clazz STRING")
      .load("data/students.txt")

    val scoreDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", ",")
      .schema("id STRING,cId STRING,score DOUBLE")
      .load("data/score.txt")

    /**
     * 1、row_number
     * 2、rank
     * 3、sum
     * 4、count
     * 5、avg
     * 6、max
     * 7、min
     * 8、lag
     * 9、lead
     */

    /**
     * row_number
     * 1、计算每个班级总分前十的学生
     */

    scoreDF
      //统计学生的总分
      .groupBy($"id")
      .agg(sum($"score") as "sumScore")
      //关联学生表获取班级
      .join(studentDF, "id")
      //增加排名字段
      .withColumn("r", row_number() over Window.partitionBy($"clazz").orderBy($"sumScore".desc))
      //取前十
      .where($"r" <= 10)
    //.show(100)

    /**
     * avg over
     * 2、计算每个学生的总分和班级平均分的差距
     */

    scoreDF
      //统计学生的总分
      .groupBy($"id")
      .agg(sum($"score") as "sumScore")
      //关联学生表获取班级
      .join(studentDF, "id")
      //增加平均分
      .withColumn("avgSumScore", avg($"sumScore") over Window.partitionBy($"clazz"))
      //计算差距
      .withColumn("score_diff", round($"sumScore" - $"avgSumScore", 2))
    //.show()

    /**
     * sum over  没有order by情况，类似分组求和
     * count over
     */

    scoreDF
      //统计学生的总分
      .groupBy($"id")
      .agg(sum($"score") as "sumScore")
      //关联学生表获取班级
      .join(studentDF, "id")
      //班级所有人的总分
      .withColumn("clazzSumSco", sum($"sumScore") over Window.partitionBy($"clazz"))
      //班级的人数
      .withColumn("num", count($"clazz") over Window.partitionBy($"clazz"))
      //平均分
      .withColumn("sumAvgSco", round($"clazzSumSco" / $"num"))
    //.show()

    /**
     * sum over如果有order by ，累计求和
     */
    scoreDF
      //统计学生的总分
      .groupBy($"id")
      .agg(sum($"score") as "sumScore")
      //关联学生表获取班级
      .join(studentDF, "id")
      //计算累计值
      .withColumn("leijiScore", sum($"sumScore") over Window.partitionBy($"clazz").orderBy($"sumScore".desc))
    //.show()


    /**
     * lag
     * 统计每个班级总分前一名和后一名差距
     *
     */

    scoreDF
      //统计学生的总分
      .groupBy($"id")
      .agg(sum($"score") as "sumScore")
      //关联学生表获取班级
      .join(studentDF, "id")
      //取前一名的总分
      .withColumn("next_score", lag("sumScore", 1, 0) over Window.partitionBy($"clazz").orderBy($"sumScore".desc))
      //计算差距
      .withColumn("score_diff", $"next_score" - $"sumScore")
      .show()
  }
}
